#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


###############################################################################
## Trifocal 2D camera location estimation from coplanar points. The
## input to this program are files produced by extract_edgels.



from __future__ import print_function
import sys
import time
import matplotlib
matplotlib.use('Agg') 
from pylab import *
import scipy.io
import scipy.interpolate

from numpy import dot
from scipy.optimize import fmin
from camori import EdgelSet, PicturePinhole, quaternion_to_matrix, measure_error, quaternion_product, dir_colors

import simplejson

import Image

set_printoptions(precision=3)


def recur(f,x):
    if x in f:
        y=f[x]
        return [y]+recur(f,y[0])
    else:
        return []

def recur2(f,x):
    if x in f:
        y=f[x]
        return [y]+recur2(f,y)
    else:
        return []

def pplot(x,*a,**kw):
    if len(x.shape)>1:
        plot(x[:,0], x[:,1], *a, **kw)
    else:
        plot(x[0], x[1], *a, **kw)


def pvec(x,v,*a,**kw):
    plot(c_[[x[0]*ones(v.shape[0]),x[0]+v[:,0]] ], 
         c_[[x[1]*ones(v.shape[0]),x[1]+v[:,1]] ], 
         *a,**kw)
        
def distinct(ia, ib):
    return not (ia[0]==ia[1] or ia[0]==ia[2] or ia[0]==ia[3] or 
                ia[1]==ia[2] or ia[1]==ia[3] or ia[2]==ia[3] or 
                ib[0]==ib[1] or ib[0]==ib[2] or ib[0]==ib[3] or 
                ib[1]==ib[2] or ib[1]==ib[3] or ib[2]==ib[3]    )

def target_f(x, M, k):
    res = np.abs(dot(M[:,:-1], x) - M[:,-1])
    q = argsort(res)
    return res[q[k]]




def plot_sol(sol):
    a = sol[:2]
    o = array([0,0])
    b = sol[2:]

    pvec(o,10* Io.obs[io], 'r--', lw=1)
    pvec(a,10* Ia.obs[ia], 'r--', lw=1)
    pvec(b,10* Ib.obs[ib], 'r--', lw=1)

    pplot(a, 'bs')
    pplot(o, 'gs')
    pplot(b, 'rs')


if __name__ == '__main__':

    ## Plot stuff immediately
    ion()

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    #PlotStuff=False
    PlotStuff=True

    #################################################################
    ## Load image and initialize pic object

    ## Sets filename from input argument
    if len(sys.argv) < 2:
        print(sys.argv[0], '<job_file.json>')
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    f_ini = job_params['first_frame']
    f_end = job_params['last_frame']

    oldgain = job_params['extraction_gain']
    gain = job_params['matching_gain']

    aa = np.load(fileroot +'/covs.npz')
    Mcov = aa['pcov']
    Sinv = inv(Mcov)

    Nsmp = 0 #number of accumulated samples

    palldelta = []
    nalldelta = []

    pcovacc = 0
    ncovacc = 0

    for ind in range(f_ini, f_end+1):
        filename = fileroot+'/edgels/edgels-%04d.npz'%(ind)
        Ia = EdgelSet(filename, oldgain, gain)

        ## Label, the direction we are using.
        lab = 1

        Ia.calc_obs(lab)
        Ia.cluster_points(lab)

        heads = (x for x in Ia.dl if not x in Ia.dr)

        for h in heads:
            end,num = Ia.dl[h]
            i = h
            while i != end:
                ii = Ia.dnl[i]

                pdelta_ = Ia.descriptors[lab][i] - Ia.descriptors[lab][ii]
                pdelta = sqrt(dot(dot(pdelta_,Sinv),pdelta_))
                
                palldelta.append(pdelta)
                pcovacc += dot(c_[pdelta],c_[pdelta].T)

                ## Add the same vector, with opposite direction
                pdelta_ = (Ia.descriptors[lab][i][[3,4,5,0,1,2]]
                          - Ia.descriptors[lab][ii][[3,4,5,0,1,2]])
                pdelta = sqrt(dot(dot(pdelta_,Sinv),pdelta_))
                palldelta.append(pdelta)
                pcovacc += dot(c_[pdelta],c_[pdelta].T)

                ## Pick some other random sample from dataset, to
                ## collect data on the "general" descriptor distances.
                ri = random_integers(Ia.descriptors[lab].shape[0])-1

                ndelta_ = Ia.descriptors[lab][i] - Ia.descriptors[lab][ri]
                ndelta = sqrt(dot(dot(ndelta_,Sinv),ndelta_))
                nalldelta.append(ndelta)
                ncovacc += dot(c_[ndelta],c_[ndelta].T)

                ndelta = (Ia.descriptors[lab][i][[3,4,5,0,1,2]]
                          - Ia.descriptors[lab][ri][[3,4,5,0,1,2]])
                ndelta = sqrt(dot(dot(ndelta_,Sinv),ndelta_))
                nalldelta.append(ndelta)
                ncovacc += dot(c_[ndelta],c_[ndelta].T)

                Nsmp += 2
                i=ii

    palldelta.append(0)
    palldelta.append(1e4)
    nalldelta.append(0)
    nalldelta.append(1e4)
    Nsmp+=2

    ## Convert to array, and sort the data
    palldelta = array(palldelta)
    palldelta = abs(palldelta)
    palldelta.sort(0)
    nalldelta = array(nalldelta)
    nalldelta = abs(nalldelta)
    nalldelta.sort(0)

    ## Calculate the covariance matrices.
    pcov = pcovacc / Nsmp
    ncov = ncovacc / Nsmp

    print(pcov)
    print(ncov)

    P = mgrid[:Nsmp+.0]/(Nsmp-1)

    # Plot stuff
    figure(1, figsize=(8,6))
    suptitle('Estimated CDF')
    ax1 = subplot(211)
    semilogx(palldelta, P, '-', lw=2)
    xlim(1e-3, 1e3)
    ylim(0,1)
    grid()

    ax2 = subplot(212, sharex=ax1, sharey=ax1)
    semilogx(nalldelta, P, '-', lw=2)
    xlim(1e-3, 1e3)
    ylim(0,1)
    grid()
    
    savefig(fileroot+'/mplt1.png', dpi=100)

    x = 10**mgrid[-2:4.0:0.01]
    figure(2, figsize=(6,6))
    suptitle('ROC curve')
    pcdf = scipy.interpolate.interp1d(palldelta, P)(x)
    ncdf = scipy.interpolate.interp1d(nalldelta, P)(x)
    plot(ncdf,pcdf, '-', lw=2)
    grid()
    axis('equal')
    axis([-0.1,1.1,-0.1,1.1])

    savefig(fileroot+'/mplt2.png', dpi=100)

    outfile = open(fileroot +'/MahaROC.npz', 'w')
    savez(outfile, pcdf=pcdf, ncdf=ncdf)
